import os
from utils.utils3 import *
import shutil
import pandas as pd
import argparse

def loop_dataset_Cubic_Spline_Interpolation(read_dir, save_dir, n=151):
    x = produce_x(n)
    for file in os.listdir(read_dir):
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        # ordinates = pd.read_csv(read_file_path).to_numpy()
        ordinates=np.loadtxt(read_file_path, delimiter=",")
        ordinates = set0(ordinates)
        ordinates_interpolation = Cubic_Spline_Interpolation(ordinates, x, n)
        ordinates_interpolation = set0(ordinates_interpolation)
        ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
        ordinates_interpolation.columns = ['x', 'y']
        ordinates_interpolation.to_csv(save_file_path, index=False)


def loop_dataset_B_Spline_Interpolation(read_dir, save_dir, n=151):
    error_count = 0
    print("以下翼型B样条无法处理，请用三次样条进行插值：")
    for file in os.listdir(read_dir):
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        x_bspine = produce_x_BSpline(2001)  # 一定要是奇数
        ordinates = pd.read_csv(read_file_path).to_numpy()

        try:
            ordinates = set0(ordinates)
            ordinates_interpolation = B_Spline_Interpolation(ordinates, x_bspine, n, file)
            ordinates_interpolation = set0(ordinates_interpolation)
            ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
            ordinates_interpolation.columns = ['x', 'y']
            ordinates_interpolation.to_csv(save_file_path, index=False)
        except:
            error_count += 1
            print("{}: error".format(error_count), file)
    print()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--read_dir", type=str, help="待处理翼型所在文件夹路径", required=True)
    parser.add_argument("-s", "--save_dir", type=str, help="处理后的翼型保存路径", required=True)
    parser.add_argument("-n", "--nums", type=int, help="翼型插值坐标点数", required=True)
    parser.add_argument('-m', '--method', choices=['bspline', 'csi'], help="插值方法", required=True)
    args = parser.parse_args()
    if args.method == "bspline":
        loop_dataset_B_Spline_Interpolation(read_dir=args.read_dir,
                                            save_dir=args.save_dir,
                                            n=args.nums)
    elif args.method == "csi":
        loop_dataset_Cubic_Spline_Interpolation(read_dir=args.read_dir,
                                                save_dir=args.save_dir,
                                                n=args.nums)
    print("第二阶段处理后的翼型已保存在{}路径".format(args.save_dir))
